import os
import cv2
import numpy as np
import yaml
from discret_image import active_size_map_discretization

root = "../datas/"
config_path = os.path.join(root, "outdoor_final.yaml")
png_path = os.path.join(root, "outdoor_final.png")
graph_path = os.path.join(root, "outdoor_final_graph.yaml")
vis_png_path = os.path.join(root, "outdoor_final_vis.png")
node_radius = 30


def save_graph_to_yaml(graph, config_path, graph_path):
    with open(config_path, "r") as f:
        config = yaml.safe_load(f)
    resolution = config["resolution"]
    x, y, yaw = config["origin"]
    yaw = np.deg2rad(yaw)
    R = np.array([[np.cos(yaw), -np.sin(yaw)], [np.sin(yaw), np.cos(yaw)]])
    T = np.array([x, y])

    # 图像坐标 → 世界坐标（带Y轴翻转）
    height, width = cv2.imread(png_path).shape[:2]
    transformed_nodes = []
    for node in graph["nodes"]:
        x_img, y_img = node[:2]
        y_img = height - y_img  # Y轴翻转
        world_coords = R @ (np.array([x_img, y_img]) * resolution) + T
        transformed_nodes.append(world_coords.tolist())
    graph["nodes"] = [[x, y, 0] for x, y in transformed_nodes]

    with open(graph_path, "w") as f:
        yaml.dump(graph, f, default_flow_style=False)


def get_graph_by_discretization(img_path, config_path, graph_path, root=[0, 0]):
    """
    通过图像离散化生成图结构，包含图像预处理步骤
    """
    # 1. 读取图像
    img = cv2.imread(img_path, cv2.IMREAD_GRAYSCALE)

    # 2. 图像预处理
    print("[预处理] 开始图像降噪和增强...")
    img_denoised = cv2.medianBlur(img, 3)
    kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (3, 3))
    img_closed = cv2.morphologyEx(img_denoised, cv2.MORPH_CLOSE, kernel, iterations=1)
    img_smooth = cv2.GaussianBlur(img_closed, (3, 3), 0)
    _, binary_img = cv2.threshold(img_smooth, 150, 255, cv2.THRESH_BINARY)

    # 3. 读取配置参数
    with open(config_path, "r") as f:
        config = yaml.safe_load(f)
    resolution = config["resolution"]
    x, y, yaw = config["origin"]
    # ROS map_saver: yaw为弧度，像素坐标原点在左下角，y轴向上
    # 图像坐标系：左上角为(0,0)，x向右，y向下
    height, width = img.shape[:2]
    # 世界坐标转像素坐标
    root_world = np.array(root[:2], dtype=np.float32)
    # 先减去origin平移，再旋转
    R = np.array(
        [[np.cos(yaw), -np.sin(yaw)], [np.sin(yaw), np.cos(yaw)]], dtype=np.float32
    )
    T = np.array([x, y], dtype=np.float32)
    # 世界坐标转到地图坐标
    map_xy = (R.T @ (root_world - T)) / resolution
    # 地图坐标转到像素坐标（y轴翻转）
    root_pixel = np.array([map_xy[0], height - 1 - map_xy[1]])
    root_pixel = np.round(root_pixel).astype(int)
    root_pixel[0] = np.clip(root_pixel[0], 0, img.shape[1] - 1)
    root_pixel[1] = np.clip(root_pixel[1], 0, img.shape[0] - 1)
    root_pixel = tuple(root_pixel)

    # 5. 使用漫水填充标记可行区域
    h, w = img.shape
    mask = np.zeros((h + 2, w + 2), dtype=np.uint8)
    flooded = binary_img.copy()
    flood_color = 100
    lo_diff, hi_diff = 10, 10
    flags = 4 | (lo_diff << 8) | cv2.FLOODFILL_FIXED_RANGE
    cv2.floodFill(flooded, mask, root_pixel, flood_color, lo_diff, hi_diff, flags)

    # 6. 调用离散化函数
    runnable_color = flood_color
    result_img, graph = active_size_map_discretization(
        flooded, runnable_color=runnable_color, node_radius=node_radius
    )

    # 7. 【新增】统计节点度数并删除孤立点
    print("[新增] 正在统计节点度数并删除孤立点...")
    nodes = graph["nodes"]
    edges = graph["edges"]

    # 7.1 统计度数
    degrees = [0] * len(nodes)
    for u, v in edges:
        degrees[u] += 1
        degrees[v] += 1

    # 7.2 找出孤立点
    isolated_nodes = set(i for i, d in enumerate(degrees) if d == 0)
    print(f"  >> 孤立点数量: {len(isolated_nodes)}")

    # 7.3 构建新节点列表和索引映射
    new_nodes = []
    old_to_new = {}
    for old_idx, node in enumerate(nodes):
        if old_idx not in isolated_nodes:
            old_to_new[old_idx] = len(new_nodes)
            new_nodes.append(node)

    # 7.4 构建新边列表
    new_edges = []
    for u, v in edges:
        if u not in isolated_nodes and v not in isolated_nodes:
            new_u = old_to_new[u]
            new_v = old_to_new[v]
            new_edges.append([new_u, new_v])

    # 7.5 更新图结构
    graph["nodes"] = new_nodes
    graph["edges"] = new_edges
    print(f"  >> 更新后节点数: {len(new_nodes)}, 边数: {len(new_edges)}")

    # 8. 可视化结果
    cv2.imwrite(vis_png_path, result_img)

    # 9. 保存图结构
    save_graph_to_yaml(graph, config_path, graph_path)
    return graph


def yml2txt(graph_path, save_dir):
    # 加载
    with open(graph_path, "r") as f:
        graph = yaml.safe_load(f)
    # 保存点坐标
    point_file_path = os.path.join(save_dir, "points.txt")
    with open(point_file_path, "w") as f:
        for node in graph["nodes"]:
            x, y, z = node
            f.write(f"{x} {y}\n")
    # 生成邻接矩阵并保存
    num_node = len(graph["nodes"])
    adj_matrix = np.eye(num_node, dtype=np.int32)
    for edge in graph["edges"]:
        node1_idx = edge[0]
        node2_idx = edge[1]
        adj_matrix[node1_idx, node2_idx] = True
        adj_matrix[node2_idx, node1_idx] = True

    adj_file_path = os.path.join(save_dir, "adj_matrix.txt")
    with open(adj_file_path, "w") as f:
        f.write(str(num_node))
        f.write("\n")
        for i in range(num_node):
            for j in range(num_node):
                s = 1 if adj_matrix[i, j] else 0
                f.write(f"{s} ")
            f.write("\n")


if __name__ == "__main__":
    get_graph_by_discretization(png_path, config_path, graph_path)
    yml2txt(graph_path, root)
